package common

import (
	"bytes"
	"compress/gzip"
	"io"
)

// 添加压缩标志
const (
	FlagCompressed = 1 << iota
)

// TunnelMessage 定义隧道通信的消息格式
type TunnelMessage struct {
	Type      MessageType `json:"type"`
	Target    string      `json:"target,omitempty"`
	RequestID string      `json:"request_id,omitempty"`
	Payload   []byte      `json:"payload,omitempty"`
	Flags     int         `json:"flags,omitempty"` // 消息标志
}

// CompressPayload 压缩消息体
func (m *TunnelMessage) CompressPayload() error {
	if len(m.Payload) < 1024 { // 小于1KB的数据不压缩
		return nil
	}

	var buf bytes.Buffer
	gw := gzip.NewWriter(&buf)
	if _, err := gw.Write(m.Payload); err != nil {
		return err
	}
	if err := gw.Close(); err != nil {
		return err
	}

	m.Payload = buf.Bytes()
	m.Flags |= FlagCompressed
	return nil
}

// DecompressPayload 解压消息体
func (m *TunnelMessage) DecompressPayload() error {
	if m.Flags&FlagCompressed == 0 {
		return nil
	}

	gr, err := gzip.NewReader(bytes.NewReader(m.Payload))
	if err != nil {
		return err
	}
	defer gr.Close()

	data, err := io.ReadAll(gr)
	if err != nil {
		return err
	}

	m.Payload = data
	m.Flags &^= FlagCompressed
	return nil
}

// MessageType 消息类型
type MessageType int

const (
	// 消息类型常量
	TypeRegister MessageType = iota
	TypeHTTPRequest
	TypeHTTPResponse
	TypeHeartbeat
	TypeError
	TypeFileChunk         // 添加文件块类型
	TypeFileComplete      // 文件传输完成类型
	TypeFileChunkRequest  // 文件分片请求类型
	TypeFileChunkResponse // 文件分片响应类型
)

// HTTPRequest 封装HTTP请求信息
type HTTPRequest struct {
	RequestID string            `json:"request_id"`
	Method    string            `json:"method"`
	Path      string            `json:"path"`
	Headers   map[string]string `json:"headers"`
	Body      []byte            `json:"body"`
	IsFile    bool              `json:"is_file,omitempty"`   // 是否是文件传输
	FileID    string            `json:"file_id,omitempty"`   // 文件ID
	FileSize  int64             `json:"file_size,omitempty"` // 文件大小
}

// HTTPResponse 封装HTTP响应信息
type HTTPResponse struct {
	RequestID  string            `json:"request_id"`
	StatusCode int               `json:"status_code"`
	Headers    map[string]string `json:"headers"`
	Body       []byte            `json:"body"`
}

// 添加文件传输相关结构
type FileChunk struct {
	FileID      string `json:"file_id"`      // 文件唯一标识
	ChunkIndex  int    `json:"chunk_index"`  // 块索引
	TotalChunks int    `json:"total_chunks"` // 总块数
	ChunkSize   int    `json:"chunk_size"`   // 块大小
	Data        []byte `json:"data"`         // 块数据
}

// FileChunkRequest 文件分片请求
type FileChunkRequest struct {
	RequestID string `json:"request_id"`
	FilePath  string `json:"file_path"`
	Offset    int64  `json:"offset"`
	Size      int64  `json:"size"`
}

// FileChunkResponse 文件分片响应
type FileChunkResponse struct {
	RequestID string `json:"request_id"`
	Data      []byte `json:"data"`
	Error     string `json:"error,omitempty"`
}
